################ Scheme Interpreter in Python

## (c) Peter Norvig, 2010; See http://norvig.com/lispy2.html

################ 270loc 3x.version1 Symbol, Procedure, Env classes

from __future__ import division
import re, sys

class Symbol(str): pass

def Sym(s, symbol_table={}):
    "Find unique entry for s in symbol table; install new one if needed."
    if s not in symbol_table: symbol_table[s] = Symbol(s)
    return symbol_table[s]

S_quote, S_if, S_set, S_define, S_lambda, S_begin, S_definemacro, = map(Sym, 
 "quote    if    set!   define    lambda    begin    define-macro".split())

S_quasiquote, S_unquote, S_unquotesplicing = map(Sym,
 "quasiquote    unquote    unquote-splicing".split())

class Procedure(object):
    "A user-defined Scheme procedure."
    def __init__(self, parms, exp, env):
        self.parms, self.exp, self.env = parms, exp, env
    def __call__(self, *args): 
        return eval(self.exp, Env(self.parms, args, self.env))

class Env(dict):
    "An environment: a dict of {'var':val} pairs, with an outer Env."
    def __init__(self, parms=(), args=(), outer=None):
        # Bind parm list to correspond args, or single parm to list of args
        self.update({parms:list(args)} if isa(parms,Symbol) else zip(parms,args))
        self.outer = outer
    def find(self, var):
        "Find the innermost Env where var does/could appear."
        if var in self: return self
        elif self.outer is None: raise LookupError(var)
        else: return self.outer.find(var)

def is_pair(x): return x != [] and isa(x, list)

def callcc(proc):
    "Call proc with current continuation"
    ball = RuntimeWarning()
    def throw(retval): ball.message = retval; raise ball
    try:
        return proc(throw)
    except RuntimeWarning as w:
        if w == ball: return ball.message
        else: raise w

def add_globals(self):
    "Add some Scheme standard procedures."
    import math, cmath, operator as op
    self.update(vars(math))
    self.update(vars(cmath))
    self.update({
     '+':op.add, '-':op.sub, '*':op.mul, '/':op.div, 'not':op.not_, 
     '>':op.gt, '<':op.lt, '>=':op.ge, '<=':op.le, '=':op.eq, 
     'equal?':op.eq, 'eq?':op.is_, 'length':len, 'cons':lambda x,y:[x]+list(y), 
     'car':lambda x:x[0], 'cdr':lambda x:x[1:], 'append':op.add,  
     'list':lambda *x:list(x), 'list?': lambda x:isa(x,list),
     'null?':lambda x:x==[], 'symbol?':lambda x: isa(x, Symbol),
     'boolean?':lambda x: isa(x, bool), 'pair?':is_pair, 
     'port?': lambda x:isa(x,file), 'apply':lambda proc,l: proc(*l), 
     'eval':lambda x: eval(expand(x)), 'load':lambda fn: load(fn), 'call/cc':callcc,
     'open-input-file':open,'close-input-port':lambda p: p.close(), 
     'open-output-file':lambda f:open(f,'w'), 'close-output-port':lambda p: p.close(),
     'write':lambda x,port=sys.stdout:port.write(to_string(x)),
     'display':lambda x,port=sys.stdout:port.write(x if isa(x,str) else to_string(x))})
    return self

isa = isinstance

global_env = add_globals(Env())

################ eval (tail recursive)

def eval(x, env=global_env):
    "Evaluate an expression in an environment."
    while True:
        if isa(x, Symbol):        # variable reference
            return env.find(x)[x]
        elif not isa(x, list):    # constant literal
            return x                
        elif x[0] is S_quote:     # (quote exp)
            (_, exp) = x
            return exp
        elif x[0] is S_if:        # (if test conseq alt)
            (_, test, conseq, alt) = x
            x = (conseq if eval(test, env) else alt)
        elif x[0] is S_set:       # (set! var exp)
            (_, var, exp) = x
            env.find(var)[var] = eval(exp, env)
            return None
        elif x[0] is S_define:    # (define var exp)
            (_, var, exp) = x
            env[var] = eval(exp, env)
            return None
        elif x[0] is S_lambda:    # (lambda (var*) exp)
            (_, vars, exp) = x
            return Procedure(vars, exp, env)
        elif x[0] is S_begin:     # (begin exp*)
            for exp in x[1:-1]:
                eval(exp, env)
            x = x[-1]
        else:                     # (proc exp*)
            exps = [eval(exp, env) for exp in x]
            proc = exps.pop(0)
            if isa(proc, Procedure):
                x = proc.exp
                env = Env(proc.parms, exps, proc.env)
            else:
                return proc(*exps)

################ parse, read, and user interaction

def read(s):
    "Read a Scheme expression from a string."
    return read_from(tokenize(s))

def parse(s):
    "Parse a program: read and expand/error-check it."
    return expand(read(s))

def tokenize(s):
    """Separate string s into tokens.  A token can be:
    a comment (which is ignored); a paren or ,@ or , or quote or quasiquote; a non-string atom;
    or a string consisting of quotes around (non-quotes or backslash plus anything)."""
    tokens = re.findall(r"""\s*(;.*|,@|[('`,)]|[^\s('");]+|"(?:[\\].|[^\\"])*")\s*""", s)
    return [t for t in tokens if not t.startswith(';')]

quotes = {"'":S_quote, "`":S_quasiquote, ",":S_unquote, ",@":S_unquotesplicing}

def read_from(tokens):
    "Read an expression from a sequence of tokens."
    if not tokens:
        raise SyntaxError('premature end of expression')
    token = tokens.pop(0)
    if '(' == token:
        L = []
        while tokens[0] != ')':
            L.append(read_from(tokens))
        tokens.pop(0) # ')'
        return L
    elif ')' == token:
        raise SyntaxError('unexpected )')
    elif token in quotes:
        return [quotes[token], read_from(tokens)]
    else:
        return atom(token)

def atom(token):
    'Numbers become numbers; #t and #f are booleans; "..." string; otherwise Symbol.'
    if token == '#t': return True
    elif token == '#f': return False
    elif token[0] == '"': return token[1:-1].replace('\\"','"').replace('\\n', '\n')
    try: return int(token)
    except ValueError:
        try: return float(token)
        except ValueError:
            try: return complex(token.replace('i','j',1))
            except ValueError:
                return Sym(token)

def to_string(x):
    "Convert a Python object back into a Lisp-readable string."
    if x is True: return "#t"
    elif x is False: return "#f"
    elif isa(x, Symbol): return x
    elif isa(x, str): return '"%s"' % x.replace('"', '\\"').replace('\n', '\\n')
    elif isa(x, list): return '('+' '.join(map(to_string, x))+')'
    elif isa(x, complex): return str(x).replace('j', 'i')
    else: return str(x)

def load(filename):
    "Eval every exp in a file."
    tokens = tokenize(open(filename).read())
    while tokens:
        eval(expand(read_from(tokens)))

def repl(prompt='lispy> '):
    "A prompt-read-eval-print loop."
    while True:
        val = eval(parse(raw_input(prompt)))
        if val is not None: print to_string(val)

################ expand

def expand(x):
    "Walk tree of x, making optimizations/fixes, and signaling SyntaxError."
    require(x, x!=[])                      # () => Error
    if not isa(x, list):                   # constant => unchanged
        return x
    elif x[0] is S_quote:                  # (quote exp)
        require(x, len(x)==2)
        return x
    elif x[0] is S_if:                     # (if t c) => (if t c None)
        if len(x)==3: x = x + [None]
        require(x, len(x)==4)
        return map(expand, x)
    elif x[0] is S_set:                    # (set! non-var exp) => Error
        var = x[1]
        require(x, len(x)==3); require(x, isa(var, Symbol), "can only set! symbol")
        return [S_set, var, expand(x[2])]
    elif x[0] in (S_define,S_definemacro): # (define non-var/list exp) => Error
        require(x, len(x)>=3)
        de, v, body = x[0], x[1], x[2:]
        if isa(v, list) and v:             # (define (f args) body)
            f, args = v[0], v[1:]          #  => (define f (lambda (args) body))
            return expand([de, f, [S_lambda, args]+body])
        else:
            require(x, len(x)==3 and isa(v, Symbol))
            exp = expand(x[2])
            if de is S_definemacro: 
                expval = eval(exp)
                require(x, callable(expval), "macro must be a procedure")
                macro_table[v] = expval
            return [S_define, v, exp]
    elif x[0] is S_begin:
        if len(x)==1: return None          # (begin) => None
        else: return map(expand, x)
    elif x[0] is S_lambda:                 # (lambda (x) e1 e2) 
        require(x, len(x)>=3)              #  => (lambda (x) (begin e1 e2))
        vars, body = x[1], x[2:]
        require(x, isa(vars, Symbol) or 
                (isa(vars, list) and all(isa(v, Symbol) for v in vars)),
                "illegal argument list")
        exp = body[0] if len(body) == 1 else [S_begin] + body
        return [S_lambda, vars, expand(exp)]   
    elif x[0] is S_quasiquote:             # `x => expand_quasiquote(x)
        require(x, len(x)==2)
        return expand_quasiquote(x[1])
    elif isa(x[0], Symbol) and x[0] in macro_table:
        return expand(macro_table[x[0]](*x[1:]))# (m arg...) => macroexpand if m in MACROS
    else:
        return map(expand, x)              # (f arg...) => expand each

def require(x, predicate, msg="wrong length"):
    "Signal a syntax error if predicate is false."
    if not predicate: raise SyntaxError(to_string(x)+': '+msg)

def expand_quasiquote(x):
    """Expand `x => 'x; `,x => x; `(,@x y) => (append x y) """
    if not is_pair(x):
        return [S_quote, x]
    require(x, x[0] is not S_unquotesplicing, "can't splice here")
    if x[0] is S_unquote:
        require(x, len(x)==2)
        return x[1]
    elif is_pair(x[0]) and x[0][0] is S_unquotesplicing:
        require(x[0], len(x[0])==2)
        return [Sym('append'), x[0][1], expand_quasiquote(x[1:])]
    else:
        return [Sym('cons'), expand_quasiquote(x[0]), expand_quasiquote(x[1:])]

macro_table = {}

eval(parse("""(begin

(define-macro and (lambda args 
   (if (null? args) #t
       (if (= (length args) 1) (car args)
           `(if ,(car args) (and ,@(cdr args)) #f)))))

;; More macros can go here

)"""))
